import cv2
import numpy as np

BLACK_THRESHOLD = 135  # 根据用户提供的数值调整
WHITE_THRESHOLD = 150  # 根据用户提供的数值调整

def process_circles(img, lower_color, upper_color):
    """
    基于颜色范围检测圆形轮廓
    :param img: 输入图像
    :param lower_color: 颜色范围下限
    :param upper_color: 颜色范围上限
    :return: 检测到的圆心坐标数组
    """
    # 创建颜色掩码
    mask = cv2.inRange(img, lower_color, upper_color)

    # 形态学操作去噪
    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    centers = []
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > 200:  # 调整面积阈值以过滤小区域
            perimeter = cv2.arcLength(contour, True)
            if perimeter > 0:
                circularity = 4 * np.pi * (area / (perimeter * perimeter))
                if circularity > 0.6:  # 调整圆形度阈值以筛选更圆的区域
                    x, y, w, h = cv2.boundingRect(contour)
                    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
                    center_x = x + int(w / 2)
                    center_y = y + int(h / 2)
                    centers.append([center_x, center_y])

    # 按Y坐标排序并取前5个点
    centers.sort(key=lambda point: point[1])
    return centers[:5]


def search_black_circle(img):
    """
    检测黑色棋子
    :param img: 输入图像
    :return: 黑色棋子的圆心坐标数组
    """
    lower_black = np.array([0, 0, 0], dtype=np.uint8)
    upper_black = np.array([BLACK_THRESHOLD, BLACK_THRESHOLD, BLACK_THRESHOLD], dtype=np.uint8)
    return process_circles(img, lower_black, upper_black)


def search_white_circle(img):
    """
    检测白色棋子
    :param img: 输入图像
    :return: 白色棋子的圆心坐标数组
    """
    lower_white = np.array([WHITE_THRESHOLD, WHITE_THRESHOLD, WHITE_THRESHOLD], dtype=np.uint8)
    upper_white = np.array([255, 255, 255], dtype=np.uint8)
    return process_circles(img, lower_white, upper_white)


def map_to_tic_tac_toe(centers, img_width, img_height):
    """
    将棋子中心点映射到3x3棋盘
    :param centers: 检测到的棋子中心坐标列表
    :param img_width: 图像宽度
    :param img_height: 图像高度
    :return: 3x3棋盘状态
    """
    grid_width = img_width // 3
    grid_height = img_height // 3
    board = [[0]*3 for _ in range(3)]  # 初始化3x3棋盘，0表示无棋子

    for center in centers:
        x, y = center

        # 确定所在的列
        col = min(2, x // grid_width)
        # 确定所在的行
        row = min(2, y // grid_height)

        # 更新棋盘状态
        board[row][col] = 1  # 假设黑棋子

    return board


if __name__ == "__main__":
    cap = cv2.VideoCapture(1)

    # 设置摄像头分辨率
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    if not cap.isOpened():
        print("Cannot open camera")
        exit()

    while True:
        ret, img_original = cap.read()
        if not ret:
            print("Can't receive frame (stream end?). Exiting ...")
            break

        # 确认分辨率是否设置成功
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        if width != 640 or height != 480:
            # 调整分辨率失败，可能需要弹出错误或使用cv2.resize调整
            print(f"原始分辨率不是640x480，当前是{width}x{height}，强制调整...")
            img_original = cv2.resize(img_original, (640, 480))
            width, height = 640, 480

        # 计算裁剪的中间部分
        crop_size_w = 320  # 裁剪宽度
        crop_size_h = 320  # 裁剪高度
        crop_x = (width - crop_size_w) // 2
        crop_y = (height - crop_size_h) // 2

        # 裁剪中间部分
        img_cropped = img_original[crop_y:crop_y + crop_size_h, crop_x:crop_x + crop_size_w]

        # 检测黑色和白色棋子
        black_array = search_black_circle(img_cropped)
        white_array = search_white_circle(img_cropped)

        # 映射到3x3棋盘
        black_board = map_to_tic_tac_toe(black_array, img_cropped.shape[1], img_cropped.shape[0])
        white_board = map_to_tic_tac_toe(white_array, img_cropped.shape[1], img_cropped.shape[0])

        # 合并棋盘状态
        tic_tac_toe_board = [[0]*3 for _ in range(3)]
        for i in range(3):
            for j in range(3):
                if black_board[i][j] == 1:
                    tic_tac_toe_board[i][j] = 1  # 黑棋子
                if white_board[i][j] == 1:
                    tic_tac_toe_board[i][j] = -1  # 白棋子

        # 打印棋盘状态
        print("Tic Tac Toe Board:")
        for row in tic_tac_toe_board:
            print(row)
        print("===")

        # 绘制裁剪框
        cv2.rectangle(img_original, (crop_x, crop_y), (crop_x + crop_size_w, crop_y + crop_size_h), (0, 0, 255), 2)

        # 显示原始图像
        cv2.imshow("Original View", img_original)



        # 按 'q' 键退出
        if cv2.waitKey(1) == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()